from fastapi import FastAPI
from pydantic import Field

from config.app_config import AppConfig
from config.base.base_nacos_config import BaseNacosConfig
from core.config import Config
from core.open_telemetry.otel_config import setup_otel


class OpenTelemetryConfig(BaseNacosConfig):
    enable: bool = Field(description='是否启用OpenTelemetry', default=False, alias='open_telemetry.enable')
    otlp_endpoint: str = Field(description='OTLP 收集器端点', default='http://180.184.146.229:4317', alias='open_telemetry.otlp_endpoint')
    enable_httpx: bool = Field(description='是否启用HTTPX自动埋点', default=True, alias='open_telemetry.enable_httpx')
    enable_pymysql: bool = Field(description='是否启用PyMySQL自动埋点', default=False, alias='open_telemetry.enable_pymysql')
    enable_aiohttp: bool = Field(description='是否启用AIOHTTP自动埋点', default=True, alias='open_telemetry.enable_aiohttp')
    enable_pymongo: bool = Field(description='是否启用PyMongo自动埋点', default=False, alias='open_telemetry.enable_pymongo')
    enable_aio_pika: bool = Field(description='是否启用AIO_PIKA自动埋点', default=True, alias='open_telemetry.enable_aio_pika')


def init_open_telemetry(app: FastAPI, app_config: AppConfig):
    open_telemetry_config = Config.get(OpenTelemetryConfig)
    if not open_telemetry_config.enable:
        return
    setup_otel(
        service_name=app_config.app_name + '_' + app_config.app_env.value,
        otlp_endpoint=open_telemetry_config.otlp_endpoint,
        app=app,
        enable_httpx=open_telemetry_config.enable_httpx,
        enable_pymysql=open_telemetry_config.enable_pymysql,
        enable_pymongo=open_telemetry_config.enable_pymongo,
        enable_aiohttp=open_telemetry_config.enable_aiohttp,
        enable_aio_pika=open_telemetry_config.enable_aio_pika,
    )
